package com.itwxz.framework.config;

import io.undertow.server.DefaultByteBufferPool;
import io.undertow.websockets.jsr.WebSocketDeploymentInfo;
import org.springframework.boot.web.embedded.undertow.UndertowServletWebServerFactory;
import org.springframework.boot.web.server.WebServerFactoryCustomizer;
import org.springframework.context.annotation.Configuration;

/**
 * UndertowConfig：用于自定义 Undertow Servlet 容器的配置。
 * 主要用于设置 WebSocket 所需的 BufferPool，以提高性能或控制内存管理方式。
 *
 * @author gastronome
 */
@Configuration
public class UndertowConfig implements WebServerFactoryCustomizer<UndertowServletWebServerFactory> {
	/**
	 * customize 方法是 WebServerFactoryCustomizer 接口提供的回调方法，
	 * 用于定制 UndertowServletWebServerFactory。
	 */
	@Override
	public void customize(UndertowServletWebServerFactory factory) {
		factory.addDeploymentInfoCustomizers(deploymentInfo -> {
			WebSocketDeploymentInfo webSocketDeploymentInfo = new WebSocketDeploymentInfo();

			// 设置自定义 Buffer Pool，false 表示非直接内存（如需性能更高可设 true）
			webSocketDeploymentInfo.setBuffers(new DefaultByteBufferPool(false, 1024));

			// 一定要以这个 key 注册
			deploymentInfo.addServletContextAttribute(
					WebSocketDeploymentInfo.ATTRIBUTE_NAME, webSocketDeploymentInfo
			);
		});
	}
}
